# encoding=utf-8
import os
import cv2
import imutils
import xml.etree.ElementTree as ET
from detector_yolov7 import Detector
from tqdm import tqdm


def autolabel_voc(image_path="", info=None):
    """
    Args:
        :param image_path : absolute path of image
        :param info:
    """
    # Read image and detect object
    image = cv2.imread(image_path)
    (image_height, image_width, image_depth) = image.shape
    # detector = Detector(weight=weight)
    # _, info = detector.detect(image)
    # Set up root node, on VOC format, root node is called annotation
    root = ET.Element("annotation")
    # Get storage folder and image name
    folder_path, image_name = os.path.split(image_path)
    image_prefix, _ = os.path.splitext(image_name)
    # Get the upper level directory
    folder_upper = os.path.dirname(folder_path)
    folder_upper += os.sep
    # Get folder name
    folder_name = folder_path.replace(folder_upper, "")
    # Set up the first child node "folder"
    folder = ET.SubElement(root, "folder")
    folder.text = folder_name
    # Set up the second child node "filename"
    folder = ET.SubElement(root, "filename")
    folder.text = image_name
    # Set up the third child node "filename"
    folder = ET.SubElement(root, "path")
    folder.text = image_path
    # Set up the fourth child node "source"
    source = ET.SubElement(root, "source")
    database = ET.SubElement(source, "database")
    database.text = "unknown"
    # Set up the fifth child node "size"
    pic_size = ET.SubElement(root, "size")
    pic_width = ET.SubElement(pic_size, "width")
    pic_width.text = str(image_width)
    pic_height = ET.SubElement(pic_size, "height")
    pic_height.text = str(image_height)
    pic_depth = ET.SubElement(pic_size, "depth")
    pic_depth.text = str(image_depth)
    # Set up the sixth child node "segmented"
    segmented = ET.SubElement(root, "segmented")
    segmented.text = "0"
    # Set every object coordinate
    # print("prediction:", prediction)
    for order in range(info['box_nums']):
        # print("det:", det)
        object_det = ET.SubElement(root, "object")
        object_name = ET.SubElement(object_det, "name")
        object_name.text = detector.names[int(info['class_ids'][order])]
        object_pose = ET.SubElement(object_det, "pose")
        object_pose.text = "Unspecified"
        object_truncated = ET.SubElement(object_det, "truncated")
        object_truncated.text = "0"
        object_difficult = ET.SubElement(object_det, "difficult")
        object_difficult.text = "0"
        object_bndbox = ET.SubElement(object_det, "bndbox")
        object_xmin = ET.SubElement(object_bndbox, "xmin")
        # print("str(int(det[0]))", str(int(det[0])))
        object_xmin.text = str(int(info['boxes'][order][0]))
        object_ymin = ET.SubElement(object_bndbox, "ymin")
        # print("str(int(det[1]))", str(int(det[1])))
        object_ymin.text = str(int(info['boxes'][order][1]))
        object_xmax = ET.SubElement(object_bndbox, "xmax")
        # print("str(int(det[2]))", str(int(det[2])))
        object_xmax.text = str(int(info['boxes'][order][2]))
        object_ymax = ET.SubElement(object_bndbox, "ymax")
        # print("str(int(det[3]))", str(int(det[3])))
        object_ymax.text = str(int(info['boxes'][order][3]))
    # Storage node tree information in ElementTree and save it as xml file
    tree = ET.ElementTree(root)
    xml_path = os.path.join(folder_upper, "Annotations", image_prefix + ".xml")
    tree.write(xml_path)


if __name__ == "__main__":
    # Set up the directory of storing pictures
    image_storage_folder = "/home/kaijia/algo-env/datasets/hq_safeguard_dataset/274A/Images"
    weight = "weights/yolov7_hq_safeguard.pt"
    detector = Detector(weight=weight, conf_thres=0.7, iou_thres=0.3)
    directory_upper = os.path.dirname(image_storage_folder)
    anno_path = os.path.join(directory_upper, "Annotations")
    if not os.path.exists(anno_path):
        os.mkdir(anno_path)
    image_list = os.listdir(image_storage_folder)
    image_count = 0
    # class_list = ["operator"]
    image_all = len(image_list)
    for img in tqdm(image_list):
        image_count += 1
        print("Processing {0} / {1}".format(image_count, image_all))
        img_abs_path = os.path.join(image_storage_folder, img)
        image = cv2.imread(img_abs_path)
        # info = detector.detect_class(image, class_names=class_list)
        _, info = detector.detect(image)
        print("Processing {0}".format(img_abs_path))
        autolabel_voc(image_path=img_abs_path, info=info)
